#include "model_landing_3dof_m.h"
/**
 *      三自由度模型,弹体速度系和弹体系重合，攻角，侧滑作为等效发动机摆角，方向遵循此方向。见航飞动
 *      m x y z vx vy vz theta theta_dot phi dphi_dot
 *      m dv/dt = P*cos(alpha)*cos(beta) - X - mgsin(theta)
 *      m v dtheta/dt = P*sin(alpha) - mgcos(theta)
 *      m v cos(theta) dphi/dt = P*cos(alpha)sin(beta) 
 *      dx/dt = V*cos(theta)cos(phi)
 *      dy/dt = V*sin(theta)
 *      dz/dt = -V*cos(theta)sin(phi)
 *      //dtheta_dot/dt = -P*sin(alpha)*rTB/m*(rG*rG)
 *      //dphi_dot/dt = -P*cos(alpha)sin(beta)*rTB/m*(rG*rG)
 *      dm/dt = -P/Ig*g 
 *      
 * 
 * 
 */

/**
 * @brief 按照给定的初始状态初始化状态和控制量，采用的方法是线性插值，终点到初始点的一条线段
 * 
 * @param X 状态量
 * @param U 输入量
 */
void model_landing_3dof_m::initialize(Eigen::Matrix<double, n_states, K> &X, Eigen::Matrix<double, n_inputs, K> &U) {

    X.setZero();
    U.setZero();
    for (size_t k = 0; k < K; k++) {
        double alpha2 = double(k)/K;
        //double alpha1 = 1. - alpha2;

        X(0,k) = (m_dry - m_wet) * alpha2 + m_wet;
        X(1,k) = (rx_final - rx_init) * alpha2 + rx_init;
        X(2,k) = (ry_final - ry_init) * alpha2 + ry_init;
        X(3,k) = (rz_final - rz_init) * alpha2 + rz_init;

        X(4,k) = (vx_final - vx_init) * alpha2 + vx_init;
        X(5,k) = (vy_final - vy_init) * alpha2 + vy_init;
        X(6,k) = (vz_final - vz_init) * alpha2 + vz_init;

        X(7,k) = (theta_final - theta_init) * alpha2 + theta_init;
        X(8,k) = (phi_final - phi_init) * alpha2 + phi_init;
        

        U(0,k) = 0.0;
        U(1,k) = 0.0;
        U(2,k) = 0.0;
    }

}
model_landing_3dof_m::StateVector model_landing_3dof_m::ode(const StateVector &x, const ControlVector &u)
{
    //         0 1 2 3 4  5  6   7     8            0   1    2
    //状态变量 m x y z vx vy vz theta  phi  控制变量 P alpha beta
    const double P = u(0, 0);
    const double alpha = u(1, 0);
    const double beta = u(2, 0);
    const double m = x(0, 0);
    const double vx = x(4, 0);
    const double vy = x(5, 0);
    const double vz = x(6, 0);
    const double theta = x(7, 0);
    const double phi = x(8, 0);

    const double X_m = C_D*(vx*vx + vy*vy + vz*vz);
    const double sigma_1 = (X_m / m) + g * sin(theta) - (P * cos(alpha) * cos(beta)) / m;
    const double sigma_2 = sqrt(vx * vx + vy * vy + vz * vz);

    StateVector f;


    f(0, 0) = -P / (I_sp * g);
    f(1, 0) = vx;
    f(2, 0) = vy;
    f(3, 0) = vz;
    f(4, 0) = -cos(phi) * cos(theta) * sigma_1;
    f(5, 0) = -sin(theta) * sigma_1;
    f(6, 0) = cos(theta) * sin(phi) * sigma_1;
    f(7, 0) = (P * sin(alpha)) / (m * sigma_2) - (g * cos(theta)) / sigma_2;
    f(8, 0) = (P * cos(alpha) * sin(beta)) / (m * cos(theta) * sigma_2);


    return f;
}
model_landing_3dof_m::StateMatrix model_landing_3dof_m::state_jacobian(const StateVector &x, const ControlVector &u)
{
    //         0 1 2 3 4  5  6   7     8            0   1    2
    //状态变量 m x y z vx vy vz theta  phi  控制变量 P alpha beta
    const double P = u(0, 0);
    const double alpha = u(1, 0);
    const double beta = u(2, 0);
    const double m = x(0, 0);
    const double vx = x(4, 0);
    const double vy = x(5, 0);
    const double vz = x(6, 0);
    const double theta = x(7, 0);
    const double phi = x(8, 0);
    

    //临时变量

    model_landing_3dof_m::StateMatrix jacobian;
    jacobian.setZero();
    const double sigma_5 = P * cos(alpha) * cos(beta);
    const double sigma_6 = vx * vx + vy * vy + vz * vz;
    const double sigma_1 = g * sin(theta) + (C_D * sigma_6) / m - sigma_5 / m;
    const double sigma_2 = (C_D * sigma_6) / (m * m) - sigma_5 / (m * m);
    const double sigma_3 = m * cos(theta) * pow(sigma_6, 1.5);
    const double sigma_4 = m * pow(sigma_6, 1.5);


    jacobian(1, 4) = 1;
    jacobian(2, 5) = 1;
    jacobian(3, 6) = 1;
    jacobian(4, 0) = cos(phi) * cos(theta) * sigma_2;
    jacobian(4, 4) = -(2 * C_D * vx * cos(phi) * cos(theta)) / m;
    jacobian(4, 5) = -(2 * C_D * vy * cos(phi) * cos(theta)) / m;
    jacobian(4, 6) = -(2 * C_D * vz * cos(phi) * cos(theta)) / m;
    jacobian(4, 7) = cos(phi) * sin(theta) * sigma_1 - g * cos(phi) * cos(theta) * cos(theta);
    jacobian(4, 8) = cos(theta) * sin(phi) * sigma_1;
    jacobian(5, 0) = sin(theta) * sigma_2;
    jacobian(5, 4) = -(2 * C_D * vx * sin(theta)) / m;
    jacobian(5, 5) = -(2 * C_D * vy * sin(theta)) / m;
    jacobian(5, 6) = -(2 * C_D * vz * sin(theta)) / m;
    jacobian(5, 7) = -cos(theta) * sigma_1 - g * cos(theta) * sin(theta);
    jacobian(6, 0) = -cos(theta) * sin(phi) * sigma_2;
    jacobian(6, 4) = (2 * C_D * vx * cos(theta) * sin(phi)) / m;
    jacobian(6, 5) = (2 * C_D * vy * cos(theta) * sin(phi)) / m;
    jacobian(6, 6) = (2 * C_D * vz * cos(theta) * sin(phi)) / m;
    jacobian(6, 7) = g * cos(theta) * cos(theta) * sin(phi) - sin(phi) * sin(theta) * sigma_1;
    jacobian(6, 8) = cos(phi) * cos(theta) * sigma_1;
    jacobian(7, 0) = -P * sin(alpha) / (m * m * sqrt(sigma_6));
    jacobian(7, 4) = (g * vx * cos(theta)) / pow(sigma_6, 1.5) - (P * vx * sin(alpha)) / sigma_4;
    jacobian(7, 5) = (g * vy * cos(theta)) / pow(sigma_6, 1.5) - (P * vy * sin(alpha)) / sigma_4;
    jacobian(7, 6) = (g * vz * cos(theta)) / pow(sigma_6, 1.5) - (P * vz * sin(alpha)) / sigma_4;
    jacobian(7, 7) = (g * sin(theta)) / sqrt(sigma_6);
    jacobian(8, 0) = -P * cos(alpha) * sin(beta) / (m * m * cos(theta) * sqrt(sigma_6));
    jacobian(8, 4) = -P * vx * cos(alpha) * sin(beta) / sigma_3;
    jacobian(8, 5) = -P * vy * cos(alpha) * sin(beta) / sigma_3;
    jacobian(8, 6) = -P * vz * cos(alpha) * sin(beta) / sigma_3;
    jacobian(8, 7) = P * cos(alpha) * sin(beta) * sin(theta) / (m * cos(theta) * cos(theta) * sqrt(sigma_6));

    return jacobian;
}

model_landing_3dof_m::ControlMatrix model_landing_3dof_m::control_jacobian(const StateVector &x, const ControlVector &u)
{
    //         0 1 2 3 4  5  6   7     8            0   1    2
    //状态变量 m x y z vx vy vz theta  phi  控制变量 P alpha beta
    const double P = u(0, 0);
    const double alpha = u(1, 0);
    const double beta = u(2, 0);
    const double m = x(0, 0);
    const double vx = x(4, 0);
    const double vy = x(5, 0);
    const double vz = x(6, 0);
    const double theta = x(7, 0);
    const double phi = x(8, 0);

    model_landing_3dof_m::ControlMatrix jacobian;
    jacobian.setZero();

    const double sigma_2 = sqrt(vx * vx + vy * vy + vz * vz);
    const double sigma_1 = m * cos(theta) * sigma_2;

    jacobian(0, 0) = -1.0 / (I_sp * g);
    jacobian(4, 0) = cos(alpha) * cos(beta) * cos(phi) * cos(theta) / m;
    jacobian(4, 1) = -P * cos(beta) * cos(phi) * sin(alpha) * cos(theta) / m;
    jacobian(4, 2) = -P * cos(alpha) * cos(phi) * sin(beta) * cos(theta) / m;
    jacobian(5, 0) = cos(alpha) * cos(beta) * sin(theta) / m;
    jacobian(5, 1) = -P * cos(beta) * sin(alpha) * sin(theta) / m;
    jacobian(5, 2) = -P * cos(alpha) * sin(beta) * sin(theta) / m;
    jacobian(6, 0) = -cos(alpha) * cos(beta) * cos(theta) * sin(phi) / m;
    jacobian(6, 1) = P * cos(beta) * sin(alpha) * cos(theta) * sin(phi) / m;
    jacobian(6, 2) = P * cos(alpha) * sin(beta) * cos(theta) * sin(phi) / m;
    jacobian(7, 0) = sin(alpha) / (m * sigma_2);
    jacobian(7, 1) = P * cos(alpha) / (m * sigma_2);
    jacobian(8, 0) = cos(alpha) * sin(beta) / sigma_1;
    jacobian(8, 1) = -P * sin(alpha) * sin(beta) / sigma_1;
    jacobian(8, 2) = P * cos(alpha) * cos(beta) / sigma_1;

    return jacobian;
}

void model_landing_3dof_m::add_application_constraints(
    optimization_problem::SecondOrderConeProgram &socp,
    Eigen::Matrix<double, n_states, K> &X0,
    Eigen::Matrix<double, n_inputs, K> &U0
) {

    auto var = [&](const string &name, const vector<size_t> &indices){ return socp.get_variable(name,indices); };
    auto param = [](double &param_value){ return optimization_problem::Parameter(&param_value); };
    
    // initial state
    socp.add_constraint( (-1.0) * var("X", {0, 0}) + param(m_wet) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {1, 0}) + param(rx_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {2, 0}) + param(ry_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {3, 0}) + param(rz_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {4, 0}) + param(vx_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {5, 0}) + param(vy_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {6, 0}) + param(vz_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {7, 0}) + param(theta_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {8, 0}) + param(phi_init) == 0.0 );


    // final state
    socp.add_constraint( (1.0) * var("X", {0, K-1}) + (-m_dry) >= 0.0 );
    socp.add_constraint( (1.0) * var("X", {1, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {2, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {3, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {4, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {5, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {6, K-1}) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {7, K-1}) + PI/2 == 0.0);

    // glide slope cone 下滑道约束
    const double tan_gamma_gs = tan(max_glide_slope_angle);
    for (size_t k = 0; k < K; ++k) {
        socp.add_constraint(
            optimization_problem::norm2({ 
                (1.0) * var("X", {2, k}),  // x
                (1.0) * var("X", {4, k})   // z
            })
            <= (1.0 / tan_gamma_gs) * var("X", {3, k})  // y
        );
    }
    // const double tan_gamma_gs = tan(max_glide_slope_angle);
    // for (size_t k = 0; k < K-1; ++k) {
    //     socp.add_constraint(
    //         optimization_problem::norm2({ 
    //             (1.0) * var("X", {0, k})
    //         })
    //         <= 1/tan_gamma_gs * var("X", {1, k}) 
    //     );
    //     //socp.add_constraint(1/tan_gamma_gs * var("X", {1, k})  + (-1.0) * var("X", {0, k}) >= 0.0);
    // }
    //地面上飞行约束
    for (size_t k = 0; k < K; ++k) {
        socp.add_constraint((1.0) * var("X", {3, k}) >= 0.0);
    }
    // TODO
    //姿态约束
    // for (size_t k = 0; k < K; ++k) {
    //     socp.add_constraint((-1.0) * var("X", {4, k}) + max_attitude_angle >= 0.0);
    //     socp.add_constraint((1.0) * var("X", {4, k}) + max_attitude_angle >= 0.0);
    // }

    for (size_t k = 0; k < K; ++k) {

        // throttle control constraints
        socp.add_constraint( ( 1.0) * var("U", {0, k}) + (-P_min) >= (0.0) );
        socp.add_constraint( (-1.0) * var("U", {0, k}) + (P_max) >= (0.0) );

        // gimbal control constraints
        if(k < 0.8 * K) {
            socp.add_constraint( ( 1.0) * var("U", {1, k}) + (max_gimbal_angle) >= (0.0) );
            socp.add_constraint( (-1.0) * var("U", {1, k}) + (max_gimbal_angle) >= (0.0) );
        } else {
            socp.add_constraint( ( 1.0) * var("U", {1, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
            socp.add_constraint( (-1.0) * var("U", {1, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
        }
        if(k < 0.8 * K) {
            socp.add_constraint( ( 1.0) * var("U", {2, k}) + (max_gimbal_angle) >= (0.0) );
            socp.add_constraint( (-1.0) * var("U", {2, k}) + (max_gimbal_angle) >= (0.0) );
        } else {
            socp.add_constraint( ( 1.0) * var("U", {2, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
            socp.add_constraint( (-1.0) * var("U", {2, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
        }
    }
}

// void model_landing_3dof::add_application_constraints(
//     optimization_problem::SecondOrderConeProgram &socp,
//     Eigen::Matrix<double, n_states, K> &X0,
//     Eigen::Matrix<double, n_inputs, K> &U0
// ) {

//     auto var = [&](const string &name, const vector<size_t> &indices){ return socp.get_variable(name,indices); };
//     auto param = [](double &param_value){ return optimization_problem::Parameter(&param_value); };
    
//     // initial state
//     socp.add_constraint( (-1.0) * var("X", {0, 0}) + param(rx_init) == 0.0 );
//     socp.add_constraint( (-1.0) * var("X", {1, 0}) + param(ry_init) == 0.0 );
//     socp.add_constraint( (-1.0) * var("X", {2, 0}) + param(vx_init) == 0.0 );
//     socp.add_constraint( (-1.0) * var("X", {3, 0}) + param(vy_init) == 0.0 );
//     socp.add_constraint( (-1.0) * var("X", {4, 0}) + param(theta_init) == 0.0 );
//     socp.add_constraint( (-1.0) * var("X", {5, 0}) + param(dtheta_init) == 0.0 );


//     // final state
//     socp.add_constraint( (1.0) * var("X", {0, K-1}) == 0.0 );
//     socp.add_constraint( (1.0) * var("X", {1, K-1}) == 0.0 );
//     socp.add_constraint( (1.0) * var("X", {2, K-1}) == 0.0 );
//     socp.add_constraint( (1.0) * var("X", {3, K-1}) == 0.0 );
//     socp.add_constraint( (1.0) * var("X", {4, K-1}) == 0.0 );
//     socp.add_constraint( (1.0) * var("X", {5, K-1}) == 0.0 );

//     // glide slope cone
//     // const double tan_gamma_gs = tan(max_glide_slope_angle);
//     // for (size_t k = 0; k < K; ++k) {
//     //     socp.add_constraint(
//     //         optimization_problem::norm2({ 
//     //             (1.0) * var("X", {2, k}),  // z
//     //             (1.0) * var("X", {3, k})   // vz
//     //         })
//     //         <= (1.0 / tan_gamma_gs) * var("X", {1, k})  // x
//     //     );
//     // }
//     const double tan_gamma_gs = tan(max_glide_slope_angle);
//     for (size_t k = 0; k < K-1; ++k) {
//         socp.add_constraint(
//             optimization_problem::norm2({ 
//                 (1.0) * var("X", {0, k})
//             })
//             <= 1/tan_gamma_gs * var("X", {1, k}) 
//         );
//         //socp.add_constraint(1/tan_gamma_gs * var("X", {1, k})  + (-1.0) * var("X", {0, k}) >= 0.0);
//     }
//     //地面上飞行约束
//     for (size_t k = 0; k < K; ++k) {
//         socp.add_constraint((1.0) * var("X", {1, k}) >= 0.0);
//     }
//     // TODO
//     //姿态约束
//     // for (size_t k = 0; k < K; ++k) {
//     //     socp.add_constraint((-1.0) * var("X", {4, k}) + max_attitude_angle >= 0.0);
//     //     socp.add_constraint((1.0) * var("X", {4, k}) + max_attitude_angle >= 0.0);
//     // }

//     for (size_t k = 0; k < K; ++k) {

//         // throttle control constraints
        // socp.add_constraint( ( 1.0) * var("U", {0, k}) + (-0.1) >= (0.0) );
        // socp.add_constraint( (-1.0) * var("U", {0, k}) + (1.0) >= (0.0) );

//         // gimbal control constraints
//         if(k < 0.8 * K) {
//             socp.add_constraint( ( 1.0) * var("U", {1, k}) + (max_gimbal_angle) >= (0.0) );
//             socp.add_constraint( (-1.0) * var("U", {1, k}) + (max_gimbal_angle) >= (0.0) );
//         } else {
//             socp.add_constraint( ( 1.0) * var("U", {1, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
//             socp.add_constraint( (-1.0) * var("U", {1, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
//         }
//     }
// }

model_landing_3dof_m::StateVector model_landing_3dof_m::get_random_state() {
    StateVector X;
    X.setRandom();
    X *= 10;
    return X;
}

model_landing_3dof_m::ControlVector model_landing_3dof_m::get_random_input() {
    ControlVector U;
    U.setRandom();
    return U;
}